import pandas as pd
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# 配置中文字体显示及负号正常显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, LSTM, Bidirectional
from numpy.random import seed

# 导入tushare库用于获取股票数据
import tushare as ts
import pandas as pd
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# 配置中文字体显示及负号正常显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, LSTM, Bidirectional
from numpy.random import seed

# 导入tushare库用于获取股票数据
def normalize_data(data):
    """
    使用MinMaxScaler对输入数据进行归一化处理，将数据范围缩放到0到1之间

    参数:
    data (np.ndarray): 待归一化的二维数组形式的数据

    返回:
    np.ndarray: 归一化后的数据，以及拟合的MinMaxScaler对象（用于后续对测试数据进行同样的归一化转换）
    """
    scaler = MinMaxScaler(feature_range=(0, 1))
    normalized_data = scaler.fit_transform(data)
    return normalized_data, scaler  # 这里补上冒号，正确结束函数定义


def split_data(sequence, n_timestamp):
    """
    将时间序列数据按照指定的时间戳窗口大小进行分割，生成用于模型训练的输入特征X和对应的目标输出y

    参数:
    sequence (np.ndarray): 一维或二维的时间序列数据（如果是二维应为 (样本数, 特征数) 的格式）
    n_timestamp (int): 时间戳窗口大小，即使用前多少个时间点的数据来预测下一个时间点

    返回:
    tuple: 包含分割后的输入特征X（三维数组格式，(样本数, 时间戳窗口大小, 特征数)）和目标输出y（二维数组格式，(样本数, 特征数)）
    """
    X, y = [], []
    for i in range(len(sequence)):
        end_ix = i + n_timestamp
        if end_ix > len(sequence) - 1:
            break
        seq_x, seq_y = sequence[i:end_ix], sequence[end_ix]
        X.append(seq_x)
        y.append(seq_y)
    return np.array(X), np.array(y)
    # 设置随机数种子，确保结果可重现
    seed(1)

    tf.random.set_seed(1)

    # 股票相关参数
    token = '572b6125c826e36619a670a9943995412179dabcfb739c5c48913f46'  # 替换为真实token
    ts_code = "601939.SH"  # 建设银行股票代码，可替换为其他感兴趣的股票代码
    start_date = "20100101"
    end_date = "20241225"
    fields = ["ts_code", "trade_date", "open", "high", "low", "close", "pre_close", "change", "pct_chg", "vol", "amount"]

    # 深度学习模型相关参数
    n_timestamp = 5  # 时间戳窗口大小
    n_epochs = 30  # 训练轮数
    model_type = 1  # 1: 单层 LSTM; 2: 多层 LSTM; 3: 双向 LSTM
        # 获取股票数据
    stock_data = get_stock_data(token, ts_code, start_date, end_date, fields)
    if stock_data is None:
        print("无法获取股票数据，程序退出。")
        exit(1)
    # 将数据存储为CSV文件，方便后续查看和使用
    stock_data.to_csv("stock_{}.csv".format(ts_code), index=False)
        # 读取数据并进行预处理
    mydata = pd.read_csv("stock_{}.csv".format(ts_code), parse_dates=["trade_date"], index_col="trade_date")[["open", "high", "low", "close"]]
    mydata = mydata.iloc[::-1]  # 翻转数据，可根据实际需求决定是否需要这步操作

    # 可视化收盘价走势
    plt.figure(figsize=(12, 5))
    plt.plot(mydata["close"], label="收盘价")
    plt.title("{}股票收盘价走势".format(ts_code))
    plt.xlabel("日期")
    plt.ylabel("价格")
    plt.legend()
    plt.show()

    # 获取股价数据部分（这里可根据需求调整取数逻辑，比如取不同时间范围等）
    df3 = mydata.reset_index().iloc[30:, :4]  # 取过去30天数据（示例，可调整），只取需要的列
    df3 = df3.dropna(how='any').reset_index(drop=True)
    df3 = df3.sort_values(by='trade_date', ascending=True)

    # 计算均线数据
    df3['5'] = df3.close.rolling(5).mean()
    df3['10'] = df3.close.rolling(10).mean()
    df3['30'] = df3.close.rolling(30).mean()

    # 可视化蜡烛图及均线
    import mplfinance as mpf
    import matplotlib.dates as mdates

    mpf.plot(df3, type='candle', style='yahoo',
             title='{}股票价格走势及均线'.format(ts_code),
             ylabel='价格',
             ylabel_lower='成交量',
             volume=True,
             figsize=(12, 8),
             savefig='stock_{}_candle.png'.format(ts_code),  # 可选择保存图片
             addplot=[mpf.make_addplot(df3['5']),
                      mpf.make_addplot(df3['10']),
                      mpf.make_addplot(df3['30'])])
        # 准备用于深度学习的数据
    training_set = df3.close.values.reshape(-1, 1)  # 取收盘价列并转换为二维数组格式
    test_set = df3.close.values[-300:].reshape(-1, 1)  # 取后300个数据作为测试集（示例，可调整）

    # 数据归一化
    training_set_scaled, scaler = normalize_data(training_set)
    testing_set_scaled = scaler.transform(test_set)

    # 数据分割
    X_train, y_train = split_data(training_set_scaled, n_timestamp)
    X_test, y_test = split_data(testing_set_scaled, n_timestamp)

    # 调整数据形状，符合LSTM模型输入要求（增加维度用于表示样本数量、时间步长和特征数量）
    X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)
    X_test = X_test.reshape(X_test.shape[0], X_test.shape[1], 1)

    # 构建模型
    model = build_model(model_type, X_train.shape[1], 1)
    model.summary()

    # 编译模型
    model.compile(optimizer=tf.keras.optimizers.Adam(0.001),
                  loss='mean_squared_error')

    # 训练模型
    history = model.fit(X_train, y_train,
                        batch_size=64,
                        epochs=n_epochs,
                        validation_data=(X_test, y_test),
                        validation_freq=1,
                        verbose=1)

    # 绘制训练和验证损失曲线
    plt.plot(history.history['loss'], label='训练损失')
    plt.plot(history.history['val_loss'], label='验证损失')
    plt.title("模型训练和验证损失曲线")
    plt.xlabel("轮数")
    plt.ylabel("损失值")
    plt.legend()
    plt.show()

    # 模型预测
    predicted_stock_price = model.predict(X_test)
    predicted_stock_price = scaler.inverse_transform(predicted_stock_price)  # 反归一化
    real_stock_price = scaler.inverse_transform(y_test)

    # 可视化真实价格和预测价格对比
    plt.figure(figsize=(12, 5))
    plt.plot(real_stock_price, label='真实股价', color='blue')
    plt.plot(predicted_stock_price, label='预测股价', color='red')
    plt.title("{}股票真实价格与预测价格对比".format(ts_code))
    plt.xlabel("时间")
    plt.ylabel("价格")
    plt.legend()
    plt.show()

   # 模型评估指标计算
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score

mse = mean_squared_error(real_stock_price, predicted_stock_price)
rmse = np.sqrt(mse)
mae = mean_absolute_error(real_stock_price, predicted_stock_price)
r2 = r2_score(real_stock_price, predicted_stock_price)

print("均方误差: {:.5f}".format(mse))
print("均方根误差: {:.5f}".format(rmse))
print("平均绝对误差: {:.5f}".format(mae))
print("R2得分: {:.5f}".format(r2))


def build_model(model_type, input_shape, output_shape):
    """
    根据指定的模型类型构建相应结构的深度学习模型（LSTM相关）

    参数:
    model_type (int): 模型类型选择，1表示单层LSTM，2表示多层LSTM，3表示双向LSTM
    input_shape (tuple): 输入数据的形状，通常为 (时间步长, 特征数量)
    output_shape (int): 输出数据的维度，即预测的目标数量（这里通常为1，预测股价）

    返回:
    tf.keras.models.Sequential: 构建好的深度学习模型
    """
    model = Sequential()
    if model_type == 1:
        # 单层 LSTM
        model.add(LSTM(units=50, activation='relu',
                       input_shape=input_shape))
        model.add(Dense(units=output_shape))
    elif model_type == 2:
        # 多层 LSTM
        model.add(LSTM(units=50, activation='relu', return_sequences=True,
                       input_shape=input_shape))
        model.add(LSTM(units=50, activation='relu'))
        model.add(Dense(1))
    elif model_type == 3:
        # 双向 LSTM
        model.add(Bidirectional(LSTM(50, activation='relu'),
                                input_shape=input_shape))
        model.add(Dense(1))
    return model
    